"""主程序入口"""
import argparse
import os
import sys
sys.stdout.reconfigure(encoding='utf-8')
sys.stderr.reconfigure(encoding='utf-8')
from pathlib import Path
PROJECT_ROOT_STR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT_STR not in sys.path:
    sys.path.insert(0, PROJECT_ROOT_STR)
    
from config.path_config import PROJECT_ROOT
from config.build_config import BuildConfig, create_default_config
from builders.engine_builder import PPQEngineBuilder

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Build TensorRT engine from ONNX with PPQ quantization')
    parser.add_argument('--onnx', type=str, help='Path to ONNX file')
    parser.add_argument('--ppq_json', type=str, help='Path to PPQ JSON file')
    parser.add_argument('--engine', type=str, help='Output path for TensorRT engine')
    parser.add_argument('--workspace', type=int, default=1<<32, help='TensorRT workspace size')
    parser.add_argument('--fp16', action='store_true', help='Enable FP16 precision')
    return parser.parse_args()

def main():
    """主函数"""
    args = parse_arguments()
    
    try:
        # 创建配置
        if args.onnx and args.ppq_json and args.engine:
            config = BuildConfig(
                onnx_path=args.onnx,
                ppq_json_path=args.ppq_json,
                engine_path=args.engine,
                workspace_size=args.workspace,
                enable_fp16=args.fp16
            )
        else:
            # 使用默认配置
            from __init__ import MODELS_DIR
            config = create_default_config(MODELS_DIR)
        
        # 构建引擎
        PPQEngineBuilder.build_int8_engine(
            config.onnx_path,
            config.ppq_json_path,
            config.engine_path,
            config.workspace_size,
            config.enable_fp16
        )
        
        print(f"\033[1;32mSuccessfully generated engine: {config.engine_path}\033[0m")
        
    except Exception as e:
        print(f"\033[1;31mError: {e}\033[0m")
        return 1
    
    return 0

if __name__ == '__main__':
    exit(main())